#include "EQHeat.h"
#include <iostream>
#include <mgl2/mgl.h>

using namespace arma;
using namespace std;

void EQHeat::set_a(double diff_coef)
{
	a = diff_coef;
}

void EQHeat::set_xlimits(double lower, double upper)
{
	xlimits[0] = lower;
	xlimits[1] = upper;
}
	
void EQHeat::set_tlimits(double lower, double upper)
{
	tlimits[0] = lower;
	tlimits[1] = upper;
}
	
void EQHeat::set_BC_left(f1v BC_left)
{
	_BC_left = BC_left;
}

void EQHeat::set_BC_right(f1v BC_right)
{
	_BC_right = BC_right;
}

void EQHeat::set_IC(f1v IC)
{
	_IC = IC;
}

void EQHeat::set_hx(double h)
{
	hx = h;
}

void EQHeat::set_ht(double tau)
{
	ht = tau;
}


unsigned int EQHeat::compute_Nx()
{
	Nx = (xlimits[1] - xlimits[0])/hx;
	return Nx;
}
	
unsigned int EQHeat::compute_Nt()
{
	Nt = (tlimits[1] - tlimits[0])/ht;
}
	
void EQHeat::init()
{
	compute_Nx();
	compute_Nt();

	u.resize(Nt+1);
	for (int i = 0; i < Nt+1; ++i)
	{
		u[i].set_size(Nx-1);
	}

	BC_L = (*_BC_left)(xlimits[0]);
	BC_R = (*_BC_right)(xlimits[1]);
	for (int i = 0; i < Nx-1; ++i)
	{
		u[0](i) = (*_IC)(xlimits[0] + (i+1)*hx);
	}
	u[0].insert_rows(0,BC_L*ones<vec>(1));
	u[0].insert_rows(Nx,BC_R*ones<vec>(1));
	d = a*a*ht/hx/hx;
	// vec tmp = zeros<vec>(Nx-1);
	// tmp(0) = 1+2*d; tmp(1) = -d;
	// A = toeplitz(tmp);
}

void EQHeat::solve()
{
	init();
	cout << "Heat equation, diffusion coefficient: a = " << a << endl;
	cout << "hx = " << hx << ", ht = " << ht << endl;
	cout << "Nx = " << Nx << ", Nt = " << Nt << endl;
	// A.print("A=");
	vec tmp_b, tmp_x, tmp_y;
	// mat L, U;
	vec U(Nx-1), L(Nx-1);
	tmp_x.set_size(Nx-1);
	tmp_y.set_size(Nx-1);
	// lu(L,U,A);
	double dd = (1+2*d); 
	U(0) = dd;
	for (int i = 1; i < Nx-1; ++i)
	{
		L(i) = -d/U(i-1);
		U(i) = dd + L(i)*d;
	}

	for (int i = 1; i < Nt+1; ++i)
	{
		tmp_b = u[i-1].rows(1,Nx-1);
		// tmp_b.print("b");
		tmp_b(0) = tmp_b(0) - d*BC_L;
		tmp_b(Nx-2) = tmp_b(Nx-2) - d*BC_R;
		tmp_y.zeros();
		tmp_y(0) = tmp_b(0);
		for (int j = 1; j < Nx-1; ++j)
		{
			tmp_y(j) = tmp_b(j) - L(j)*tmp_y(j-1);
		}
		tmp_x.zeros();
		tmp_x(Nx-2) = tmp_y(Nx-2)/U(Nx-2);
		for (int j = Nx-3; j >=0; --j)
		{

			tmp_x(j) = (tmp_y(j) + d*tmp_x(j+1))/U(j);

		}

		u[i] = tmp_x;
		u[i].insert_rows(0,BC_L*ones<vec>(1));
		u[i].insert_rows(Nx,BC_R*ones<vec>(1));
	}

}

void EQHeat::save_data(const char* filename) const
{
	using namespace std;

	ofstream fout;
	fout.open(filename);
	//fout.precision(4);
	//fout.width(5);
	for (int i = 0; i < Nt+1; ++i)
	{
		fout << "# t = " << tlimits[0] + i*ht << endl;
		for (int j = 0; j < Nx; ++j)
		{
			fout << u[i](j) << ", ";			
		}
		fout << u[i](Nx) << ";" << endl << endl;
	}
}

void EQHeat::draw_data(const char* filename, std::vector<double> time) const
{
	const char* lineScheme[10] = {"p-1", "r|1", "g;1", "b:1", "cj1", "mi1", "y=1", "l1", "e1", "k1"};

	mglData mdata(Nx+1), t(Nx+1);
	for (int i = 0; i < Nx+1; ++i)
	{
		t.a[i] = tlimits[0] + i*hx;
	}
	double ymin = 0, ymax = 0;
	// for (int i = 0; i < time.size(); ++i)
	// {
	// 	if()
	// }
	mglGraph gr;
	gr.Alpha(true);
	gr.Light(true);
	for (int i = 0; i < time.size(); ++i)
	{
		int ind = (time[i]-tlimits[0])/ht;
		if (u[ind].max() > ymax)
			ymax = u[ind].max();
		if (u[ind].min() < ymin)
			ymin = u[ind].min();
	}
	gr.SetRanges(xlimits[0],xlimits[1], floor(ymin), ceil(ymax));
	// gr.Box();
	gr.Axis();
	gr.Grid("xy", "W");
	gr.Label('x',"x",1);
	gr.Label('y',"u(x,t)",1);
	char legend[256];
	double *a = new double[Nx+1];
	for (int i = 0; i < time.size(); ++i)
	{
		int ind = (time[i]-tlimits[0])/ht;
		for (int j = 0; j < Nx+1; ++j)
		{			
			a[j] = u[ind](j);
		}
		mdata.Set(a, Nx+1);
		sprintf(legend, "t = %f", time[i]);
		gr.Plot(t, mdata, lineScheme[i]);
		gr.AddLegend(legend,  lineScheme[i]);
	}

	gr.Legend(3,"A#");
	gr.WriteFrame(filename);
	delete []a;
}